# -*- coding: utf-8 -*-

"""
Created on 03/23/2022
configs.
@author: Kang Xiatao (kangxiatao@gmail.com)
"""

import os
import time
import argparse
from utils import mail_log
from easydict import EasyDict as edict
import logging
from tensorboardX import SummaryWriter


def init_config():
    parser = argparse.ArgumentParser()
    ## Training Hyperparameters ##
    parser.add_argument('--config', type=str, default='cifar10/vgg19',
                        help='config (dataset/model(network depth)) - e.g.[cifar10/vgg19, mnist/lenet5]')
    parser.add_argument('--pretrained', type=str, default='666',
                        help='path - runs/anoi/finetune_cifar10_vgg19_l2_best.pth.tar')
    parser.add_argument('--run', type=str, default='test_exp', help='experimental notes (default: test_exp)')
    parser.add_argument('--epoch', type=int, default=666)
    parser.add_argument('--batch_size', type=int, default=666)
    parser.add_argument('--l2', type=str, default=666)
    parser.add_argument('--lr_mode', type=str, default='cosine', help='cosine or preset')
    parser.add_argument('--optim_mode', type=str, default='SGD', help='SGD or Adam')
    parser.add_argument('--storage_mask', type=int, default=0)  # storage mask
    parser.add_argument('--debug', type=str, default='p')  # Debug flags (printing data, plotting, etc.)
    parser.add_argument('--dp', type=str, default='../Data', help='dataset path')

    ## Distributed Config ##
    parser.add_argument('-n', '--nodes', default=1, type=int, metavar='N')
    parser.add_argument('-g', '--gpus', default=1, type=int, help='number of gpus per node')
    parser.add_argument('-nr', '--nr', default=0, type=int, help='ranking within the nodes')

    args = parser.parse_args()

    ## Default Config ##
    exp_set = args.config.split('/')  # dataset model prune_ratio
    exp_name = f'{exp_set[0]}_{exp_set[1]}'.replace('.', '_')
    base_config = {'network': "vgg", 'depth': 19, 'dataset': 'cifar10',
                   'batch_size': 128, 'epoch': 180, 'learning_rate': 0.1, 'weight_decay': 5e-4}
    config = edict(base_config)
    config.dataset = exp_set[0]
    config.network = ''.join(list(filter(str.isalpha, exp_set[1])))
    config.depth = int(''.join(list(filter(str.isdigit, exp_set[1]))))
    if 'mnist' in config.dataset:
        config.batch_size = 256
        config.epoch = 80
        config.weight_decay = 1e-4
        config.classe = 10
    elif 'cifar' in config.dataset:
        config.batch_size = 128
        config.epoch = 180
        config.weight_decay = 5e-4
        if config.dataset == 'cifar10':
            config.classe = 10
        else:
            config.classe = 100
    elif 'imagenet' in config.dataset:
        config.batch_size = 128
        config.epoch = 300
        config.classe = 200
        if 'vgg' in config.network:
            config.weight_decay = 5e-4
        elif 'resnet' in config.network:
            config.weight_decay = 1e-4

    ## Experiment Name and Out Path ##
    summn = [exp_name]
    chekn = [exp_name]
    if len(args.run) > 0:
        summn.append(args.run)
        chekn.append(args.run)
        config.exp_name = exp_name + '_' + args.run
    summn.append("summary/")
    chekn.append("checkpoint/")
    summary_dir = ["./runs"] + summn
    ckpt_dir = ["./runs"] + chekn
    config.summary_dir = os.path.join(*summary_dir)
    config.checkpoint_dir = os.path.join(*ckpt_dir)
    print("=> config.summary_dir:    %s" % config.summary_dir)
    print("=> config.checkpoint_dir: %s" % config.checkpoint_dir)

    ## Console Parameters ##
    if args.pretrained != '666':
        config.pretrained = args.pretrained
        print("use pre-trained mode:{}".format(config.pretrained))
    else:
        config.pretrained = None
    if args.epoch != 666:
        config.epoch = args.epoch
        print("set new epoch:{}".format(config.epoch))
    if args.batch_size != 666:
        config.batch_size = args.batch_size
        print("set new batch_size:{}".format(config.batch_size))
    if args.l2 != 666:
        config.weight_decay = args.l2
        print("set new weight_decay:{}".format(config.weight_decay))
    config.lr_mode = args.lr_mode
    config.optim_mode = args.optim_mode
    config.debug = args.debug
    config.dp = args.dp

    config.send_mail_head = (config.exp_name + ' -> ' + args.run + '\n')
    config.send_mail_str = (mail_log.get_words() + '\n')
    config.send_mail_str += "=> 我能在河边钓一整天的🐟 <=\n"

    config.nodes = args.nodes
    config.gpus = args.gpus
    config.nr = args.nr

    return config


def makedirs(filename):
    if not os.path.exists(os.path.dirname(filename)):
        os.makedirs(os.path.dirname(filename))


def get_logger(name, logpath, displaying=True, saving=True):
    logger = logging.getLogger(name)
    logger.setLevel(logging.INFO)
    log_path = logpath + name + time.strftime("-%Y%m%d-%H%M%S")
    makedirs(log_path)
    if saving:
        info_file_handler = logging.FileHandler(log_path, encoding='utf-8')
        info_file_handler.setLevel(logging.INFO)
        logger.addHandler(info_file_handler)
    if displaying:
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        logger.addHandler(console_handler)

    return logger


def init_logger(config):
    makedirs(config.summary_dir)
    makedirs(config.checkpoint_dir)

    # set logger
    logger = get_logger('log', logpath=config.summary_dir + '/')
    logger.info(dict(config))
    writer = SummaryWriter(config.summary_dir)
    return logger, writer